import argparse
import json
from utils.data_generator import ENoseDataGenerator

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='电子鼻数据生成工具')
    parser.add_argument('--channels', type=int, default=4,
                       help='传感器通道数量 (默认: 4)')
    parser.add_argument('--samples', type=int, default=100,
                       help='采样点数量 (默认: 100)')
    parser.add_argument('--rate', type=float, default=10.0,
                       help='采样率(Hz) (默认: 10.0)')
    parser.add_argument('--batch', type=int, default=1,
                       help='生成的样本组数 (默认: 1)')
    parser.add_argument('--no-drift', action='store_true',
                       help='不添加信号漂移')
    parser.add_argument('--no-noise', action='store_true',
                       help='不添加噪声')
    parser.add_argument('--noise-level', type=float, default=0.1,
                       help='噪声水平 (0-1之间) (默认: 0.1)')
    parser.add_argument('--output', type=str, default='sensor_data.json',
                       help='输出文件名 (默认: sensor_data.json)')
    parser.add_argument('--format', type=str, choices=['json', 'csv', 'excel'],
                       help='输出格式 (默认: 根据文件扩展名)')
    parser.add_argument('--no-metadata', action='store_true',
                       help='不包含元数据')
    
    args = parser.parse_args()
    
    # 根据format参数设置输出文件扩展名
    if args.format:
        base_name = args.output.rsplit('.', 1)[0]
        if args.format == 'json':
            args.output = f"{base_name}.json"
        elif args.format == 'csv':
            args.output = f"{base_name}.csv"
        elif args.format == 'excel':
            args.output = f"{base_name}.xlsx"
    
    # 创建数据生成器
    generator = ENoseDataGenerator(
        num_channels=args.channels,
        sampling_rate=args.rate,
        noise_level=args.noise_level
    )
    
    # 生成数据
    data = generator.generate_batch_samples(
        num_samples=args.samples,
        batch_size=args.batch,
        add_drift=not args.no_drift,
        add_noise=not args.no_noise
    )
    
    # 保存数据到文件
    generator.save_to_file(
        data=data,
        filename=args.output,
        include_metadata=not args.no_metadata
    )
    
    print(f"已生成{args.batch}组数据，每组包含{args.channels}个通道，"
          f"每个通道{args.samples}个采样点")
    print(f"数据已保存到: {args.output}")

if __name__ == '__main__':
    main() 